import os
import psutil
import sys
sys.path.append("./")
import subprocess
import time
# from slice_time import get_slice_time
class MPSSet:
    def _get_pid(self,name):
        pids = psutil.process_iter()
        result = None
        for pid in pids:
            if pid.name() == name:
                result =  pid.pid
                break
        return result
    def set_mps(self,mps_num,event=None):
        try:
            stream = os.popen('sudo nvidia-cuda-mps-control\n', mode='w', buffering=-1) # create UI controller
            pid = self._get_pid("nvidia-cuda-mps-server") # get server pid
            if pid:
                stream.write("set_active_thread_percentage "+str(pid)+" "+str(int(mps_num))+"\n")
                stream.close()
                #print("=====close mps setting")
                if event is not None:
                    event.set()
                #print("=====close mps setting",event.value)
        except Exception as e:
            print("error happens when setting MPS as ",mps_num,e)
def close():
    """
    close MPS
    Returns:

    """
    try:
        stream = os.popen('sudo nvidia-cuda-mps-control\n', mode='w', buffering=-1)
        # time.sleep(3)
        # nvidia-smi
        pids = psutil.process_iter()
        target_pid = None
        for pid in pids:
            if pid.name() == "nvidia-cuda-mps-server":
                target_pid =  pid.pid
                break
        if target_pid:
            stream.write("quit\n")
            stream.close()
    except Exception as e:
        print("close MPS error ", e)
#close()


